//      motif_2nd_model.cc
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#include "gene_em.h"

/*
Motif2ndModel::prob_t Motif2ndModel::compute_full_probability(const Motif2ndModel::datum_t& x)
{
  prob_t p = create_prob_map(sz(x));

  double background = compute_background(x);

  BOOST_FOREACH(prob_t::value_type& i, p) {
    if (i.first.rbs_inside == 0) {
      double mix_prob = log(param_.rbs_prob[0]);
      i.second = LogNumber::mul(mix_prob, background);
    }
    else {
      int pos = i.first.rbs_position;
      double mix_prob = LogNumber::mul(log(param_.rbs_prob[1]), log(param_.rbs_position_prob[pos]));

      i.second = LogNumber::mul(mix_prob, compute_motif(x, pos, background));
    }
  }

  return p;
}

void Motif2ndModel::accumulate(const datum_t& x, const prob_t& p)
{
  vd_t count = arma::zeros<vd_t>(4);
  BOOST_FOREACH(char ch, x) {
    count[ACGT::c(ch)]++;
  }

  BOOST_FOREACH(const prob_t::value_type& i, p) {
    int rbs_inside = i.first.rbs_inside;
    int rbs_position = i.first.rbs_position;
    double posterior = i.second;

    if (rbs_inside == 0) {
      double s = accumulate_temp_.rbs_prob[0];
      accumulate_temp_.rbs_prob[0] = LogNumber::add(s, posterior);
      
      for (int c = 0; c < 4; c++) {
				double s = accumulate_temp_.background_freq[c];
				accumulate_temp_.background_freq[c] = LogNumber::add(s, log(count[c])+posterior);
      }
    }
    else {
      double s = accumulate_temp_.rbs_prob[1];
      accumulate_temp_.rbs_prob[1] = LogNumber::add(s, posterior);

      vd_t count_motif = arma::zeros<vd_t>(4);
      for (int j = rbs_position; j < rbs_position+motif_len() && j < sz(x); j++) {
				count_motif[ACGT::c(x[j])]++;
			}
			for (int c = 0; c < 4; c++) {
				double s = accumulate_temp_.background_freq[c];
				accumulate_temp_.background_freq[c] = LogNumber::add(s, log(count[c]-count_motif[c])+posterior);
      }

      s = accumulate_temp_.rbs_position_prob[rbs_position];
      accumulate_temp_.rbs_position_prob[rbs_position] = LogNumber::add(s, posterior);

      for (int j = rbs_position; j < rbs_position+motif_len() && j < sz(x); j++) {
				int pos = j-rbs_position;
				char ch = x[j];
				int c = ACGT::c(ch);
				
				double s = accumulate_temp_.position_freq[pos][c];
				accumulate_temp_.position_freq[pos][c] = LogNumber::add(s, posterior);
      }
    }
  }

  //std::cerr << "  accumulate=\n";  accumulate_temp_.print();
}

void Motif2ndModel::begin_accumulate()
{
  accumulate_temp_ = param_;
  
  BOOST_FOREACH(double& p, accumulate_temp_.rbs_prob) {
    p = LogNumber::zero();
  }

  BOOST_FOREACH(double& p, accumulate_temp_.rbs_position_prob) {
    p = LogNumber::zero();
  }

  BOOST_FOREACH(double& p, accumulate_temp_.background_freq) {
    p = LogNumber::zero();
  }

  BOOST_FOREACH(vd_t& v, accumulate_temp_.position_freq) {
    BOOST_FOREACH(double& p, v) {
      p = LogNumber::zero();
    }
  }

  //std::cerr << "  begin accumulate=\n";  accumulate_temp_.print();
}

void Motif2ndModel::end_accumulate()
{
  LogNumber::normalize(accumulate_temp_.rbs_prob);
  LogNumber::normalize(accumulate_temp_.rbs_position_prob);
  LogNumber::normalize(accumulate_temp_.background_freq);
  BOOST_FOREACH(vd_t& v, accumulate_temp_.position_freq) {
    LogNumber::normalize(v);
  }

  param_.rbs_prob = arma::exp(accumulate_temp_.rbs_prob);
  param_.rbs_position_prob = arma::exp(accumulate_temp_.rbs_position_prob);
  param_.background_freq = arma::exp(accumulate_temp_.background_freq);
  int motif_len = sz(param_.position_freq);
  for (int pos = 0; pos < motif_len; pos++) {
    param_.position_freq[pos] = arma::exp(accumulate_temp_.position_freq[pos]);
  }  

  //std::cerr << "  end accumulate=\n";  accumulate_temp_.print();
}

Motif2ndModel::prob_t Motif2ndModel::create_prob_map(int len)
{
  prob_t p;
  hidden_t h;

  h.rbs_inside = 0;
  p[h] = LogNumber::zero();

  h.rbs_inside = 1;
  for (int pos = 0; pos < len; pos++) {
    h.rbs_position = pos;
    p[h] = LogNumber::zero();
  }
  
  return p;
}

double Motif2ndModel::compute_background(const Motif2ndModel::datum_t& x)
{
  double p = LogNumber::one();
  BOOST_FOREACH(char ch, x) {
    int c = ACGT::c(ch);
    p = LogNumber::mul(p, log(param_.background_freq[c]));
  }
  return p;
}

double Motif2ndModel::compute_motif(const Motif2ndModel::datum_t& x, int pos, double background)
{
  double p = background;
  int motif_len = sz(param_.position_freq);
  if (pos + motif_len > sz(x)) return LogNumber::zero();

  for (int j = pos; j < pos + motif_len && j < sz(x); j++) {
    int c = ACGT::c(x[j]);
    p = LogNumber::div(p, log(param_.background_freq[c]));
    p = LogNumber::mul(p, log(param_.position_freq[j-pos][c]));
  }
  return p;
}

void Motif2ndModel::random_initialize(int motif_len, int seq_len)
{
  param_.rbs_prob = arma::randu<vd_t>(2);
  param_.rbs_position_prob = arma::randu<vd_t>(seq_len);
  param_.background_freq = arma::randu<vd_t>(4);
  param_.position_freq = vvd_t(motif_len);
  BOOST_FOREACH(vd_t& v, param_.position_freq) {
    v = arma::randu<vd_t>(4);
  }

  param_.rbs_prob /= arma::sum(param_.rbs_prob);
  for (int pos = seq_len-motif_len+1; pos < seq_len; pos++)
    param_.rbs_position_prob[pos] = 0;
  param_.rbs_position_prob /= arma::sum(param_.rbs_position_prob);
  param_.background_freq /= arma::sum(param_.background_freq);
  BOOST_FOREACH(vd_t& v, param_.position_freq) {
    v /= arma::sum(v);
  }
}
*/
// print functions
std::string Motif2ndModel::Parameters::toString() const
{
  std::stringstream ss;
	
  rbs_prob.print(ss, "rbs_prob=");
  rbs_position_prob.print(ss, "position_prob=");
  typedef std::map<std::string, vd_t>::value_type value_type;
  BOOST_FOREACH(const value_type& v, background_markov_freq) {
    v.second.print(ss, v.first+" --> ");
  }

  ss << "pos_freq=" << std::endl;
  for (int i = 0; i < sz(position_freq); i++) {
    char s[10];
    sprintf(s, "%d=", i);
    position_freq[i].print(ss, s);
  }
  return ss.str();
}

void Motif2ndModel::Parameters::save(const std::string& outputFile) const
{
  std::ofstream out;
  FileAccess::open2write(out, outputFile);
  
  out << "Non-RBS\t RBS" << std::endl;
  out << rbs_prob[0] << "\t " << rbs_prob[1] << std::endl;
  
  typedef std::map<std::string, vd_t>::value_type value_type;
  out << "Background" << std::endl << "Prefix\t A\t C\t G\t T" << std::endl;
  BOOST_FOREACH(const value_type& v, background_markov_freq) {
    out << v.first;
    for (int c = 0; c < 4; c++) 
      out << "\t " << v.second[c];
    out << ENDL;
  }
  
  out << "Motif" << std::endl << "Pos";
  for (int pos = 0; pos < sz(position_freq); pos++)
    out << "\t " << pos+1;
  out << std::endl;

  for (int c = 0; c < 4; c++) {
    out << ACGT::ch(c);
    for (int pos = 0; pos < sz(position_freq); pos++)
      out << "\t " << (position_freq[pos][c] == arma::max(position_freq[pos]) ? "+" : "") << position_freq[pos][c];
    out << std::endl;
  }
  
  out << "Position" << std::endl;
  for (int pos = 0; pos < sz(rbs_position_prob); pos++)
    out << pos-sz(rbs_position_prob) << "\t ";
  out << std::endl;
  for (int pos = 0; pos < sz(rbs_position_prob); pos++)
    out << rbs_position_prob[pos] << "\t ";
  out << std::endl;
  
  out.close();
}

std::string Motif2ndModel::toString() const 
{ 
	std::stringstream ss;
	ss << "model=" << std::endl << param_.toString(); 
	return ss.str();
}
